/*
 * subgradient_parallel.h
 *
 *  Created on: Sep 14, 2011
 *      Author: tqlong
 */

#ifndef SUBGRADIENT_PARALLEL_H_
#define SUBGRADIENT_PARALLEL_H_
#include "simple_gm.h"

struct NodeEdge : public std::pair<Node, Edge>
{
  NodeEdge(Node i, Edge e) : std::pair<Node, Edge>(i, e) {}
  struct NodeEdgeCompare
  {
    bool operator() (const NodeEdge& n1, const NodeEdge& n2)
    {
      return (n1.first < n2.first) ||
             (n1.first == n2.first && Edge::EdgeCompare()(n1.second, n2.second));
    }
  };
};

/** Subgradient descent solver for Linear Programming relaxation
 *  of the MAP problem
 *    max_x \sum \theta_i(x_i) + \sum \theta_ij (x_i, x_j)
 *  Equivalent :
 *    max_m \sum <\theta_i,m_i> + \sum <theta_ij, m_ij>
 *    where m is a point in the marginal polytope
 *      m_i(x_i) = P(X_i = x_i)
 *      m_ij(x_i, x_j) = P(X_i = x_i, X_j = x_j)
 *
 *  LP relaxation
 *    min_\lambda L(\lambda) =
 *      \sum_{ij\in E} s_ij(1/d_i \theta_i + \lambda_i^{ij},
 *                          1/d_j \theta_j + \lambda_j^{ij},
 *                          \theta_ij)
 *    s.t.          \sum_{j\in N_i} \lamda_i^{ij} = 0
 *  where
 *    s_ij(w_i, w_j, theta_ij) = max_eta <w_i,eta_i> + <w_j, eta_j> + <theta_ij, eta_ij>
 *    s.t eta_i, eta_j, eta_ij \in marginal polytope of the edge ij
 */
class SubGradGM : public SimpleGM
{
public:
  typedef std::map<Node, arma::vec>                    n_potential_type;
  typedef std::map<Edge, arma::mat, Edge::EdgeCompare> e_potential_type;
  typedef std::map<int, int>                           imap_type;
  typedef imap_type                                    assignment_type;
  typedef std::vector<Edge>                            evec_type;
  typedef std::vector<arma::vec>                       vvec_type;
  typedef std::map<Node, evec_type>                    adjacent_edge_type;
  typedef std::map<NodeEdge, arma::vec, NodeEdge::NodeEdgeCompare> duplicate_variable_type;

  typedef const n_potential_type::value_type&          const_pair_node_t;
  typedef n_potential_type::value_type&                pair_node_t;
  typedef const e_potential_type::value_type&          const_pair_edge_t;
  typedef e_potential_type::value_type&                pair_edge_t;
  typedef const adjacent_edge_type::value_type&        const_pair_adje_t;
  typedef adjacent_edge_type::value_type&              pair_adj_t;
  typedef duplicate_variable_type::value_type&         pair_nodeedge_t;
  typedef struct {
    /** variables for node marginals */
    n_potential_type mu_;
    /** variables for edge marginals */
    e_potential_type nue_;
    /** duplicate variables for node marginals*/
    duplicate_variable_type nu_;
    /** the Lagrangian multipliers */
    duplicate_variable_type lambda_;
    /** the learning rate */
  } optim_var_type;
public:
  SubGradGM() {}
  SubGradGM(const SimpleGM&);
  SubGradGM reduce(const assignment_type& a)
  {
    return SimpleGM::reduce(a);
  }

  /** The LP relaxation solver */
  assignment_type subgradient_lp_map();
public:
  // Optimization variables
  optim_var_type var_, tmp_;
  duplicate_variable_type d_;
  double eta_, val_;
protected:
  void solve_slave_problems(optim_var_type& var);
  void init_variables();
  void build_slave_problem(optim_var_type& var, Node i, Node j, Edge e, arma::vec& wi, arma::vec& wj);
  void solve_slave_problem(const arma::vec& wi, const arma::vec& wj, const arma::mat& phi,
                           arma::vec& nui, arma::vec& nuj, arma::mat& nue);
  void update_mu(optim_var_type& var, double& s);
  void update_lambda();
  double objective(optim_var_type& var);
  void do_line_search(double f0, double s);
 };

#endif /* SUBGRADIENT_PARALLEL_H_ */
